from helper import OptimizedModel # using the helper from <URL>
import tensorflow as tf
physical_devices = tf.config.list_physical_devices('GPU')
tf.config.experimental.set_memory_growth(physical_devices[0], True)
# 简单的计算来测试 GPU
# tf.debugging.set_log_device_placement(True)

from image_utils import create_vidio_image_loader
import time
import cv2

class NSFW:

    def __init__(self,modelPath,batch=10):
        self.fn_load_image = create_vidio_image_loader()
        self.model = OptimizedModel(modelPath)
        self.batch = batch

    def predictVidio(self,vidio_path):
        sample_rate = 1  # 每隔100帧采样一次
        # 打开视频文件
        cap = cv2.VideoCapture(vidio_path)
        frame_count = 0
        batch = 20
        images = []
        totalNum = 0
        while True:
            ret, frame = cap.read()
            if not ret:
                break
            frame_count += 1
            if frame_count % sample_rate == 0:
                images.append(frame)
                if len(images) < batch:
                    continue
                # 处理采样帧
                images_input = self.fn_load_image(images)
                predictions = self.model.predict(images_input)
                print(predictions)
                totalNum += len(images)
                images = []
        # if len(images) > 0:
        #     images_input = self.fn_load_image(images)
        #     predictions = self.sess.run(self.model.predictions, feed_dict={
        #         self.model.input: images_input})
        #     print(predictions)
        #     totalNum += len(images)
        print("frame num:{}".format(totalNum))


if __name__ == "__main__":
    nsfw = NSFW('./data/savedModelFormat_FP32',10)
    # nsfw.savedModelFormat("./data/savedModelFormat")
    total_start_time = time.time()
    data_result = nsfw.predictVidio("./img/big-buck-bunny_trailer.mp4")
    print(time.time()-total_start_time)
